package zeromq

import (
	"encoding/json"
	"github.com/honwee/awm/internal/model"
	"github.com/honwee/awm/internal/repository"
	"github.com/honwee/awm/pkg/config"
	"github.com/honwee/awm/pkg/logger"
	"github.com/pebbe/zmq4"
	"go.uber.org/zap"
	"time"
)

type Server struct {
	socket          *zmq4.Socket
	traceRepo       *repository.TraceRepository
	workflowRunRepo *repository.WorkflowRunRepository
	log             *zap.SugaredLogger
}

func NewServer(traceRepo *repository.TraceRepository, workflowRunRepo *repository.WorkflowRunRepository) (*Server, error) {
	socket, err := zmq4.NewSocket(zmq4.PULL)
	if err != nil {
		return nil, err
	}

	return &Server{
		socket:          socket,
		traceRepo:       traceRepo,
		workflowRunRepo: workflowRunRepo,
		log:             logger.Logger.Sugar(),
	}, nil
}

func (s *Server) Start() error {
	addr := config.GlobalConfig.ZeroMQ.LogReceiverAddress
	if err := s.socket.Bind(addr); err != nil {
		return err
	}
	s.log.Infof("ZeroMQ server started on %s", addr)

	go s.receiveLoop()
	return nil
}

func (s *Server) Stop() error {
	return s.socket.Close()
}

func (s *Server) receiveLoop() {
	for {
		msg, err := s.socket.RecvBytes(0)
		if err != nil {
			s.log.Errorf("Failed to receive message: %v", err)
			continue
		}

		var trace model.Trace
		if err := json.Unmarshal(msg, &trace); err != nil {
			s.log.Errorf("Failed to unmarshal message: %v", err)
			continue
		}

		// 处理工作流运行记录
		if err := s.handleWorkflowRun(&trace); err != nil {
			s.log.Errorf("Failed to handle workflow run: %v", err)
			continue
		}

		// 保存追踪日志
		if err := s.traceRepo.Create(&trace); err != nil {
			s.log.Errorf("Failed to save trace: %v", err)
			continue
		}

		s.log.Debugf("Received and saved trace: %s, workflow_run_id: %s", trace.Event, trace.WorkflowRunID)
	}
}

type WorkflowData struct {
	Inputs struct {
		UserID string `json:"sys.user_id"`
		AppID  string `json:"sys.app_id"`
	} `json:"inputs"`
	CreatedAt int64 `json:"created_at"`
}

type WorkflowFinishedData struct {
	Status     string `json:"status"`
	FinishedAt int64  `json:"finished_at"`
}

func (s *Server) handleWorkflowRun(trace *model.Trace) error {
	switch trace.Event {
	case "workflow_started":
		// 创建工作流运行记录
		var data WorkflowData
		if err := json.Unmarshal(trace.Data, &data); err != nil {
			return err
		}

		run := &model.WorkflowRun{
			WorkflowRunID: trace.WorkflowRunID,
			WorkflowID:    trace.WorkflowID,
			UserID:        data.Inputs.UserID,
			Status:        model.StatusRunning,
			StartedAt:     time.Unix(data.CreatedAt, 0),
		}
		return s.workflowRunRepo.Create(run)

	case "workflow_finished":
		// 更新工作流运行状态
		var data WorkflowFinishedData
		if err := json.Unmarshal(trace.Data, &data); err != nil {
			return err
		}

		var status string
		if data.Status == "succeeded" {
			status = model.StatusSucceeded
		} else {
			status = model.StatusFailed
		}

		finishedAt := time.Unix(data.FinishedAt, 0)
		return s.workflowRunRepo.UpdateStatusWithFinishedAt(trace.WorkflowRunID, status, &finishedAt)
	}

	return nil
}
